'''
作为预测代码，可用于对一下模型的预测
1. FPM_tripletloss_train.py: FPM_tripletloss
2. two_tripletloss_oversample_train.py:tripletloss_oversample
3. tripletloss_weightloss_train.py:tripletloss_weightloss
'''

import os
import pandas as pd
import numpy as np

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from two_triplet_resnet import resnet50

from sklearn.metrics import classification_report, confusion_matrix, roc_curve, auc



BatchSize = 16
save_path = '/home/zmy/pytorch_code/checkpoint/two_tripletloss_oversample_0.1_20210522_115229'
weight_name = 'two_tripletloss_oversample_0.1-65-1.5573547601699829-regular.pth'

weights_path = save_path+'/'+weight_name


# 自定义数据集和数据预处理
class MyDataset(Dataset):

    def __init__(self, datalist):
        self.data_info = datalist


    def __len__(self):
        return len(self.data_info)

    def __getitem__(self, item):
        patientID = self.data_info[item][0]
        label = self.data_info[item][1]
        ct_path = self.data_info[item][2]
        pet_path = self.data_info[item][3]
        pet_slope = self.data_info[item][4]
        pet_intercept = self.data_info[item][5]

        ct = np.load(ct_path)
        pet = np.load(pet_path)

        # pet图像转化HU值
        if pet_slope != 1:
            pet = pet * pet.astype(np.float64)
            pet = pet.astype(np.int32)
        pet += np.int32(pet_intercept)

        # pet图像归一化
        pet = MaxMinNormalizer(pet)

        # ct和pet进行合并
        img = merge_CT_PET(ct, pet)

        return {'image': torch.from_numpy(img), 'label': torch.tensor(label)}


# 读取文件列表
def read_csv(data_sets):
    sets_path = '/data1/zmy/data2021/origin_data/divide_csv/five/'

    # 读取数据集
    data_features = []

    for set in data_sets:
        train_data = pd.read_csv(sets_path+set)
        for j in range(len(train_data)):

            # 读取文件地址
            patientid = train_data['patientID'][j]

            ct_path = '/data1/zmy/data2021/origin_data/Slice/'+str(patientid)+'/CTSlice/'

            name_list = os.listdir(ct_path)

            pet_path = '/data1/zmy/data2021/origin_data/Slice/'+str(patientid)+'/PETSlice/'
            pet_slope = float(train_data['pet_slope'][j])
            pet_intercept = float(train_data['pet_intercept'][j])

            for it in name_list:

                cancer_type = int(train_data['cancer_type'][j]) - 1
                if cancer_type == 0 or cancer_type == 1:
                    one_feature = [patientid, cancer_type, ct_path+it, pet_path+it, pet_slope, pet_intercept]

                    data_features.append(one_feature)


    return data_features


# 图像归一化
def MaxMinNormalizer(data):
    data_max = np.max(data)
    data_min = np.min(data)
    data_normalize = 1 - (data - data_min) / (data_max - data_min)
    return data_normalize


# CT和PET融合
def merge_CT_PET(ct_array, pet_array):

    img = np.asarray([ct_array, pet_array], dtype=np.float)
    return img



# 创建网络
def resnet():
    # 创建resnet50网络
    net = resnet50()
    net.conv1 = nn.Conv2d(2, 64, kernel_size=7, stride=2, padding=3, bias=False)

    # 打印网络结构和参数量
    # print(net)
    # print("Total number of paramerters in networks is {}  ".format(sum(x.numel() for x in net.parameters())))
    return net


# 结果预测
def predict():
    

    # 创建网络结构
    net = resnet().to(device)

    # 加载测试数据
    test_data_sets = ['test.csv']
    test_list = read_csv(test_data_sets)
    test_dataset = MyDataset(test_list)
    testloader = DataLoader(test_dataset, batch_size=BatchSize, shuffle=False, num_workers=2)

    # 加载模型权重
    net.load_state_dict(torch.load(weights_path))
    net.eval()

    # 返回的预测结果及对应标签
    y_preds = []
    y_preds_ones = []  # 预测为腺癌的概率
    y_true = []

    with torch.no_grad():
        for n_iter, data in enumerate(testloader):
            print('iteration:{}\ttotal {} iterations'.format(n_iter + 1, len(testloader)))

            images = data['image'].type(torch.FloatTensor).to(device)
            labels = data['label'].to(device)

            out1, outputs = net(images)

            outputs = torch.softmax(outputs, dim=1)

            _, preds = outputs.max(1)


            # print(len(outputs.tolist()))

            result = outputs.tolist()
            for mm in range(len(result)):
                y_preds_ones.append(result[mm][1])

            y_preds.extend(preds.tolist())
            y_true.extend(labels.tolist())

            print(preds.tolist())
            print(labels.tolist())
        print(y_preds_ones)
        print(y_true)
        # 保存测试集结果（概率值）
        np.save('/home/zmy/pytorch_code/two_result/two_tripletloss_oversample_labels.npy', y_true)
        np.save('/home/zmy/pytorch_code/two_result/two_tripletloss_oversample_preds.npy', y_preds_ones)

    return y_preds, y_true


# 结果评估
def evaluation(preds, labels):

    # 计算评估指标
    target_names = ['1', '2']
    result_statis = classification_report(y_true=labels, y_pred=preds, target_names=target_names)
    print(result_statis)

    # 计算混淆矩阵
    confusion = confusion_matrix(y_true=labels, y_pred=preds)
    print(confusion)


if __name__ == '__main__':

    device = torch.device("cuda:2")

    # 预测结果
    y_preds, y_true = predict()

    print('weight name: ', weight_name)

    # 结果评估
    evaluation(y_preds, y_true)


